package com.subl.mybatis.config;

import com.subl.mybatis.xa.IssXADataSource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.BeansException;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.scheduling.annotation.Scheduled;

import javax.sql.XAConnection;
import javax.transaction.xa.XAException;
import javax.transaction.xa.XAResource;
import javax.transaction.xa.Xid;
import java.sql.SQLException;
import java.util.*;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * XA事务监控组件。定时10s执行一次[XA RECOVER]命令，判断数据库中是否有执行失败的XA事务。
 * 如果连续两次查询都查询到相同的XA事务id，也就意味中此事务执行时间超过了10s则在所有数据源
 * 中回滚此事务。
 */

public class XAMonitor implements ApplicationContextAware {

    private static final Logger logger = LoggerFactory.getLogger(XAMonitor.class);

    private ApplicationContext applicationContext;
    private Map<String, IssXADataSource> xaDataSourceMap;
    private Map<String, Map<String, Xid>> preXids;

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
        xaDataSourceMap = applicationContext.getBeansOfType(IssXADataSource.class);
        preXids = new HashMap<>();
        logger.info("XA事务监控器已经启动");
    }

    @Scheduled(fixedRate = 10, timeUnit = TimeUnit.SECONDS)
    public void checkXaTransaction() {
        logger.debug("开始执行定时检查XA超时事务...");
        Iterator<Map.Entry<String, IssXADataSource>> iterator = xaDataSourceMap.entrySet().iterator();
        Map<String, Map<String, Xid>> curXids = new HashMap<>();
        while (iterator.hasNext()) {
            XAConnection xaConnection = null;
            try {
                Map.Entry<String, IssXADataSource> next = iterator.next();
                IssXADataSource xaDataSource = next.getValue();
                xaConnection = xaDataSource.getXAConnection();
                //取当前连接中XA事务列表
                Map<String, Xid> xidInfo = getXidInfo(xaConnection);
                //和上次XA事务比较是否超时事务
                Map<String, Xid> preXidInfo = preXids.get(next.getKey());
                if (preXidInfo != null) {
                    //回滚超时事务
                    compareAndRollback(xaConnection, preXidInfo, xidInfo);
                    //重新取当前数据源中XA事务列表
                    xidInfo = getXidInfo(xaConnection);
                }
                curXids.put(next.getKey(), xidInfo);

            } catch (Exception e) {
                e.printStackTrace();
            } finally {
                try {
                    xaConnection.close();
                } catch (SQLException e) {
                    throw new RuntimeException(e);
                }
            }

        }
        preXids = curXids;
    }

    private Map<String, Xid> getXidInfo(XAConnection xaConnection) throws Exception {

        XAResource xaResource = xaConnection.getXAResource();
        Xid[] xids = xaResource.recover(XAResource.TMSTARTRSCAN);
        Map<String, Xid> xidMap = Stream.of(xids)
                .collect(Collectors.toMap(
                        xid->new String(xid.getGlobalTransactionId()) +
                                new String(xid.getBranchQualifier()),
                        xid->xid));
        return xidMap;
    }

    private void compareAndRollback(XAConnection xaConnection, Map<String, Xid> preXids, Map<String, Xid> curXids) {
        preXids.entrySet().stream()
                .filter(e->curXids.containsKey(e.getKey()))
                .forEach(e->{
                    try {
                        logger.info("存在超时XA事务：{}，执行回滚操作。", e.getKey());
                        XAResource xaResource = xaConnection.getXAResource();
                        xaResource.rollback(e.getValue());
                    } catch (SQLException | XAException ex) {
                        logger.error("事务回滚操作异常", ex);
                    }
                });
    }



}
